#include "LocalTrajExtraction.h"
#include <MMM/Motion/Plugin/ModelPosePlugin/ModelPoseSensor.h>
#include <QFileDialog>
#include <Eigen/Geometry>
#include <Eigen/Dense>
#include <MMM/Motion/Plugin/KinematicPlugin/KinematicSensor.h>
#include <MMM/Motion/Plugin/ModelPosePlugin/ModelPoseSensor.h>
#include <VirtualRobot/RobotConfig.h>
#include <VirtualRobot/RobotNodeSet.h>
#include <MMMSimoxTools/MMMSimoxTools.h>
#include <SimoxUtility/math/convert.h>

using namespace MMM;

LocalTrajExtraction::LocalTrajExtraction(MMM::MotionList motions_in, float minTimestep_in, float maxTimestep_in, int framesPerSecond_in, std::string modelNameA_in, std::string modelNameB_in) {
    this->motions = motions_in;
    this->minTimestep = minTimestep_in;
    this->maxTimestep = maxTimestep_in;
    this->framesPerSecond = framesPerSecond_in;
    this->modelNameA = modelNameA_in;
    this->modelNameB = modelNameB_in;
}

MotionPtr LocalTrajExtraction::getLocalPoseTraj() {

        std::string motion_name = "localTraj_" + modelNameB + "_inFrameOf_" + modelNameA;
        ModelPoseSensorPtr sensor (new ModelPoseSensor());

        // determine if hands should be considered
        std::string specsModelA;
        std::string specsModelB;
        if (modelNameA == "MMMHandL" || modelNameA == "MMMHandR"){
            specsModelA = modelNameA;
            modelNameA = "mmm";
        }
        if (modelNameB == "MMMHandL" || modelNameB == "MMMHandR"){
            specsModelB = modelNameB;
            modelNameB = "mmm";;
        }

        // select required motion
        MotionPtr motionA = nullptr;
        MotionPtr motionB = nullptr;
        for(auto & motion : motions)
        {
            if (!motion->getModel()) continue;
            if (std::filesystem::path(motion->getModel()->getFilename()).stem() == modelNameA){
                motionA = motion;
            }
            if (std::filesystem::path(motion->getModel()->getFilename()).stem() == modelNameB){
                motionB = motion;
            }
        }
        if(!motionA || !motionB){
            MMM_ERROR << "Required motions are missing" << std::endl;
        }

        // initialize robots and sensor
        auto poseSensorA = motionA->getSensorByType<ModelPoseSensor>("ModelPose");
        auto poseSensorB = motionB->getSensorByType<ModelPoseSensor>("ModelPose");
        VirtualRobot::RobotPtr robotA = MMM::SimoxTools::buildModel(motionA->getModel(), false);
        VirtualRobot::RobotPtr robotB = MMM::SimoxTools::buildModel(motionB->getModel(), false);
        KinematicSensorPtr kinematicSensorA = nullptr;
        VirtualRobot::RobotNodeSetPtr robotNodeSetA = nullptr;
        KinematicSensorPtr kinematicSensorB = nullptr;
        VirtualRobot::RobotNodeSetPtr robotNodeSetB = nullptr;

        if (motionA->hasSensor("ModelPose") && motionA->hasSensor("Kinematic")) {
            kinematicSensorA = motionA->getSensorByType<KinematicSensor>("Kinematic");
            robotNodeSetA = VirtualRobot::RobotNodeSet::createRobotNodeSet(robotA, "RobotNodeSet", kinematicSensorA->getJointNames(), "", "", true);
        }
        if (motionB->hasSensor("ModelPose") && motionB->hasSensor("Kinematic")) {
            kinematicSensorB = motionB->getSensorByType<KinematicSensor>("Kinematic");
            robotNodeSetB = VirtualRobot::RobotNodeSet::createRobotNodeSet(robotB, "RobotNodeSet", kinematicSensorB->getJointNames(), "", "", true);
        }

        // initialiize goal motion
        MotionPtr new_motion (new Motion(motion_name, motionB->getModel(false), motionB->getModel(), motionB->getModelProcessor(), motionB->getOriginFilePath()));

        Eigen::Matrix4f TrafoA = Eigen::Matrix4f::Identity();
        Eigen::Matrix4f TrafoB;
        Eigen::Matrix4f TrafoAtoB;
        for(float t = minTimestep; t <= maxTimestep; t += 1.0/framesPerSecond)
        {
            auto dataA = boost::dynamic_pointer_cast<ModelPoseSensorMeasurement>(poseSensorA->getMeasurement(t));
            robotA->setGlobalPose(dataA->getRootPose());
            if (robotNodeSetA) {
                robotNodeSetA->setJointValues(kinematicSensorA->getDerivedMeasurement(t)->getJointAngles());
            }

            if (specsModelA == "MMMHandL"){
                TrafoA = robotA->getRobotNode("LWsegment_joint")->getGlobalPose();
            }
            else if (specsModelA == "MMMHandR"){
                TrafoA = robotA->getRobotNode("RWsegment_joint")->getGlobalPose();
            }
            else{
                 Eigen::Matrix4f transf(Eigen::Matrix4f::Identity());
                 auto m = motionA->getModel()->getModelNode("root_joint");
                 if (m) {
                     if (m->hasMarker("LocalFrame")) {
                         for (auto marker : m->markers) {
                             if (marker->name == "LocalFrame") {
                                 transf = marker->localTransform;
                                 transf.block(0, 3, 3, 1) *= 1000.0f; // m -> mm
                             }
                         }
                     }
                 }
                 TrafoA = dataA->getRootPose() * transf;
            }


            auto dataB = boost::dynamic_pointer_cast<ModelPoseSensorMeasurement>(poseSensorB->getMeasurement(t));
            robotB->setGlobalPose(dataB->getRootPose());
            if (robotNodeSetB) {
            robotNodeSetB->setJointValues(kinematicSensorB->getDerivedMeasurement(t)->getJointAngles());}

            if (specsModelB == "MMMHandL"){
                TrafoB = robotB->getRobotNode("LWsegment_joint")->getGlobalPose();
            }
            else if (specsModelB == "MMMHandR"){
                TrafoB = robotB->getRobotNode("RWsegment_joint")->getGlobalPose();
            }
            else{
                Eigen::Matrix4f transf(Eigen::Matrix4f::Identity());
                auto m = motionB->getModel()->getModelNode("root_joint");
                if (m) {
                    if (m->hasMarker("LocalFrame")) {
                        for (auto marker : m->markers) {
                            if (marker->name == "LocalFrame") {
                                transf = marker->localTransform;
                                transf.block(0, 3, 3, 1) *= 1000.0f; // m -> mm
                            }
                        }
                    }
                }
                TrafoB = dataB->getRootPose() * transf;
            }

            // calculate local pose
            // alternativ: using RobotNodes for all object and Mat4f pose = a->toLocalPose(b->getGlobalPose)
            // gives b in fram of a
            TrafoAtoB = TrafoA.inverse()*TrafoB;
            ModelPoseSensorMeasurementPtr result_data (new ModelPoseSensorMeasurement(t, TrafoAtoB));
            sensor->addSensorMeasurement(result_data);
        }


        new_motion->addSensor(sensor);
        MMM_INFO << "Returning motion" << std::endl;
        return new_motion;
}



